from open_ai_model import getOpenAiModel

# 为了支持异步调用
import asyncio
from langchain_core.output_parsers import JsonOutputParser

model = getOpenAiModel()
chain = model | JsonOutputParser()


# 异步调用
async def async_stream():
    async for text in chain.astream(
        "以JSON格式输出法国、西班牙和日本的国家及其人口列表。"
        '使用一个带有"tountries"外部键的字典，其中包含国家列表。'
        "每个国家都应该有键'name'和'population'"
    ):
        print(text, flush=True)


# 运行异步流处理
asyncio.run(async_stream())
